import duckdb
from pandasai import SmartDataframe
from pandasai.llm.openai import OpenAI


class MultiTableAnalyzer:
    def __init__(self, file_paths, llm_instance):
        self.file_paths = file_paths
        self.conn = duckdb.connect()
        self.tables = {}
        self.load_tables()
        self.llm = llm_instance

    def load_tables(self):
        """加载所有 CSV 文件为 DuckDB 表"""
        for idx, path in enumerate(self.file_paths):
            table_name = f"table_{idx + 1}"
            self.conn.execute(f"CREATE TABLE {table_name} AS SELECT * FROM '{path}'")
            self.tables[table_name] = self.conn.execute(f"SELECT * FROM {table_name}").df()
        self.conn.close()

    def analyze_table(self, table_name, question):
        """对指定表执行自然语言分析"""
        sdf = SmartDataframe(self.tables[table_name], config={"llm": self.llm})
        return sdf.chat(question)

    def analyze_all_tables(self, question_template="请分析这张表的关键信息"):
        """统一分析所有表"""
        results = {}
        for table_name in self.tables:
            results[table_name] = self.analyze_table(table_name, question_template)
        return results


if __name__ == '__main__':
    llm = OpenAI(
        api_token="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")
    # 初始化并分析多个表
    analyzer = MultiTableAnalyzer([
        'test_cases/test.csv',
        'test_cases/melb_data.csv'
    ], llm)

    # 分析每个表
    results = analyzer.analyze_all_tables("用中文描述数据分布和关键特征")

    # 打印结果
    for table, result in results.items():
        print(f"\n【{table} 的分析结果】:\n{result}")
